import React, { useState, useEffect, useRef } from 'react';
import { Input, Button, Select, Spin, Typography, Divider, Alert, Space } from 'antd';
import { SendOutlined, ClearOutlined } from '@ant-design/icons';
import ReactMarkdown from 'react-markdown';
import { Prism as SyntaxHighlighter } from 'react-syntax-highlighter';
import { vscDarkPlus } from 'react-syntax-highlighter/dist/esm/styles/prism';
import { ollamaService, ChatMessage, ModelInfo } from '../services/ollamaService';

const { TextArea } = Input;
const { Option } = Select;
const { Title, Text } = Typography;

export default function ChatPage() {
  const [messages, setMessages] = useState<ChatMessage[]>([]);
  const [input, setInput] = useState('');
  const [selectedModel, setSelectedModel] = useState('');
  const [models, setModels] = useState<ModelInfo[]>([]);
  const [loading, setLoading] = useState(false);
  const [isStreaming, setIsStreaming] = useState(false);
  const [error, setError] = useState<string | null>(null);
  
  const messagesEndRef = useRef<HTMLDivElement>(null);
  const inputRef = useRef<HTMLTextAreaElement>(null);
  
  // 获取可用模型
  useEffect(() => {
    fetchModels();
  }, []);
  
  // 消息变化时滚动到底部
  useEffect(() => {
    scrollToBottom();
  }, [messages]);
  
  const fetchModels = async () => {
    try {
      setError(null);
      const modelList = await ollamaService.getModels();
      setModels(modelList);
      if (modelList.length > 0 && !selectedModel) {
        setSelectedModel(modelList[0].name);
      }
    } catch (err) {
      setError('获取模型失败。请检查您的Ollama服务器连接。');
      console.error('获取模型错误:', err);
    }
  };
  
  const scrollToBottom = () => {
    messagesEndRef.current?.scrollIntoView({ behavior: 'smooth' });
  };
  
  const handleSend = async () => {
    if (!input.trim() || !selectedModel) return;
    
    // 添加用户消息
    const userMessage: ChatMessage = { role: 'user', content: input };
    setMessages(prev => [...prev, userMessage]);
    setInput('');
    
    try {
      setError(null);
      setLoading(true);
      setIsStreaming(true);
      
      // 添加一个空的助手消息，将在流式响应期间填充
      setMessages(prev => [...prev, { role: 'assistant', content: '' }]);
      
      // 开始流式传输
      await ollamaService.streamChat(
        {
          model: selectedModel,
          messages: [...messages, userMessage],
          stream: true
        },
        (chunk) => {
          // 用新内容更新最后一条消息
          setMessages(prev => {
            const newMessages = [...prev];
            const lastMessageIndex = newMessages.length - 1;
            newMessages[lastMessageIndex] = {
              ...newMessages[lastMessageIndex],
              content: newMessages[lastMessageIndex].content + chunk
            };
            return newMessages;
          });
        },
        () => {
          setIsStreaming(false);
          setLoading(false);
          if (inputRef.current) {
            inputRef.current.focus();
          }
        }
      );
    } catch (err) {
      setError('发送消息出错。请检查与Ollama的连接。');
      console.error('发送消息错误:', err);
      setLoading(false);
      setIsStreaming(false);
    }
  };
  
  const handleClear = () => {
    setMessages([]);
  };
  
  const renderMessage = (message: ChatMessage, index: number) => {
    const isUser = message.role === 'user';
    
    return (
      <div 
        key={index} 
        style={{ 
          marginBottom: 24,
          display: 'flex',
          flexDirection: 'column',
          backgroundColor: isUser ? '#f9f9f9' : '#fff',
          borderRadius: 8,
          padding: 16,
        }}
      >
        <div style={{ marginBottom: 8 }}>
          <Text strong style={{ color: isUser ? '#333' : '#10a37f' }}>
            {isUser ? '你' : 'AI'}
          </Text>
        </div>
        
        <div>
          {isUser ? (
            <div style={{ whiteSpace: 'pre-wrap' }}>{message.content}</div>
          ) : (
            <ReactMarkdown
              components={{
                code({ node, inline, className, children, ...props }) {
                  const match = /language-(\w+)/.exec(className || '');
                  return !inline && match ? (
                    <SyntaxHighlighter
                      style={vscDarkPlus}
                      language={match[1]}
                      PreTag="div"
                      {...props}
                    >
                      {String(children).replace(/\n$/, '')}
                    </SyntaxHighlighter>
                  ) : (
                    <code className={className} {...props}>
                      {children}
                    </code>
                  );
                }
              }}
            >
              {message.content}
            </ReactMarkdown>
          )}
        </div>
      </div>
    );
  };
  
  return (
    <div className="chat-container">
      <div style={{ flex: 1, overflowY: 'auto', paddingRight: 10 }}>
        {messages.length === 0 ? (
          <div style={{ 
            display: 'flex', 
            flexDirection: 'column', 
            alignItems: 'center', 
            justifyContent: 'center',
            height: '50vh', 
            textAlign: 'center',
            color: '#888'
          }}>
            <Title level={2} style={{ color: '#10a37f' }}>丝路文通开发部AI</Title>
            <Text>选择一个模型并开始与您的  AI 模型聊天</Text>
          </div>
        ) : (
          messages.map(renderMessage)
        )}
        <div ref={messagesEndRef} />
      </div>
      
      {error && (
        <Alert 
          message="错误" 
          description={error} 
          type="error" 
          showIcon 
          closable 
          onClose={() => setError(null)} 
          style={{ marginBottom: 16 }}
        />
      )}
      
      <div style={{ marginTop: 16 }}>
        <Space style={{ marginBottom: 16 }}>
          <Select
            value={selectedModel}
            onChange={setSelectedModel}
            style={{ width: 200 }}
            placeholder="选择模型"
            disabled={loading}
          >
            {models.map(model => (
              <Option key={model.name} value={model.name}>{model.name}</Option>
            ))}
          </Select>
          
          <Button 
            type="default" 
            icon={<ClearOutlined />} 
            onClick={handleClear}
            disabled={messages.length === 0 || loading}
          >
            清空聊天
          </Button>
        </Space>
        
        <div style={{ display: 'flex' }}>
          <TextArea
            ref={inputRef}
            value={input}
            onChange={e => setInput(e.target.value)}
            placeholder="输入信息..."
            autoSize={{ minRows: 2, maxRows: 6 }}
            onPressEnter={(e) => {
              if (!e.shiftKey) {
                e.preventDefault();
                handleSend();
              }
            }}
            disabled={loading}
            style={{ flex: 1, marginRight: 16 }}
          />
          <Button
            type="primary"
            icon={<SendOutlined />}
            onClick={handleSend}
            disabled={!input.trim() || loading || !selectedModel}
            loading={loading}
            style={{ height: 'auto' }}
          >
            发送
          </Button>
        </div>
        <Text type="secondary" style={{ display: 'block', marginTop: 8 }}>
          按 Shift+Enter 换行，按 Enter 发送
        </Text>
      </div>
    </div>
  );
} 